import numpy as np
from scipy.spatial.distance import pdist, squareform
import matplotlib.pyplot as plt

class DensityPeakCluster:
    def __init__(self, n_clusters=None, dc=None, method='gaussian', auto_select_dc=True, percent=2.0):
        self.n_clusters = n_clusters      # 聚类数量
        self.dc = dc                      # 截断距离
        self.method = method              # 密度计算方法（'cutoff'或'gaussian'）
        self.auto_select_dc = auto_select_dc  # 是否自动选择dc
        self.percent = percent            # 自动选择dc时的百分比
        self.rho = None                   # 局部密度
        self.delta = None                 # 相对距离
        self.nearest_higher = None        # 最近更高密度点
        self.labels = None                # 聚类标签
        self.centers = None               # 聚类中心索引
        self.distances = None             # 距离矩阵

    def fit(self, X):
        # 计算距离矩阵
        self.distances = squareform(pdist(X, 'euclidean'))
        
        # 自动选择截断距离dc
        if self.auto_select_dc or self.dc is None:
            self.dc = self._select_dc()
        
        # 计算局部密度rho
        self.rho = self._compute_density()
        
        # 计算相对距离delta和最近更高密度点
        self.delta, self.nearest_higher = self._compute_delta()
        
        # 选择聚类中心
        if self.n_clusters is None:
            self.n_clusters = self._suggest_clusters()
        else:
            self.centers = self._select_centers()
        
        # 分配聚类标签
        self.labels = self._assign_clusters()
        
        return self

    def _select_dc(self):
        """自动选择dc使得平均邻居数占总点数的percent%"""
        triu = np.triu(self.distances, k=1)
        distances = triu[triu > 0]
        sorted_dists = np.sort(distances)
        index = int(len(sorted_dists) * self.percent / 100)
        return sorted_dists[index]

    def _compute_density(self):
        """计算局部密度"""
        n = self.distances.shape[0]
        rho = np.zeros(n)
        if self.method == 'cutoff':
            for i in range(n):
                rho[i] = np.sum(self.distances[i] < self.dc) - 1  # 排除自身
        elif self.method == 'gaussian':
            for i in range(n):
                rho[i] = np.sum(np.exp(-(self.distances[i]/self.dc)**2)) - 1
        return rho

    def _compute_delta(self):
        """计算相对距离delta和最近更高密度点"""
        n = self.distances.shape[0]
        delta = np.full(n, np.inf)
        nearest_higher = np.full(n, -1, dtype=int)
        
        # 按密度降序排序索引
        order = np.argsort(-self.rho)
        
        for i in range(n):
            idx = order[i]
            # 处理密度最高的点
            if i == 0:
                delta[idx] = np.max(self.distances[idx])
                continue
            # 查找更高密度的点
            higher_rho_indices = order[:i]
            min_dist = np.min(self.distances[idx, higher_rho_indices])
            min_idx = higher_rho_indices[np.argmin(self.distances[idx, higher_rho_indices])]
            delta[idx] = min_dist
            nearest_higher[idx] = min_idx
        return delta, nearest_higher

    def _select_centers(self):
        """根据gamma值选择聚类中心"""
        gamma = self.rho * self.delta
        return np.argsort(-gamma)[:self.n_clusters]

    def _suggest_clusters(self):
        """通过决策图建议聚类数量（需用户手动观察）"""
        plt.figure(figsize=(10, 6))
        plt.scatter(self.rho, self.delta, s=40)
        plt.xlabel('Density (rho)')
        plt.ylabel('Distance (delta)')
        plt.title('Decision Graph')
        plt.show()
        raise ValueError("请通过决策图手动选择聚类中心，并设置n_clusters参数")

    def _assign_clusters(self):
        """分配聚类标签"""
        n = self.distances.shape[0]
        labels = -np.ones(n, dtype=int)
        
        # 标记聚类中心
        for cluster_id, center_idx in enumerate(self.centers):
            labels[center_idx] = cluster_id
        
        # 按密度降序分配标签
        order = np.argsort(-self.rho)
        for idx in order:
            if labels[idx] == -1:
                labels[idx] = labels[self.nearest_higher[idx]]
        return labels

    def plot_decision_graph(self):
        """绘制决策图"""
        plt.figure(figsize=(10, 6))
        plt.scatter(self.rho, self.delta, c=self.labels, cmap='viridis', s=40)
        plt.xlabel('Density (rho)')
        plt.ylabel('Distance (delta)')
        plt.title('Decision Graph with Cluster Colors')
        plt.colorbar()
        plt.show()
    
    def plot_results(self, X):
        """聚类结果可视化"""
        plt.figure(figsize=(12, 5))
        
        # 原始数据
        plt.subplot(121)
        plt.scatter(X[:, 0], X[:, 1], c='gray', s=30)
        plt.title("Original Data")
        
        # 聚类结果
        plt.subplot(122)
        # 区分噪声点（标签为-1）和聚类点
        noise_mask = self.labels == -1
        valid_mask = ~noise_mask
        
        plt.scatter(X[valid_mask, 0], X[valid_mask, 1], 
                    c=self.labels[valid_mask], cmap='viridis', s=50)
        plt.scatter(X[self.centers, 0], X[self.centers, 1],
                    s=200, marker='*', c='red', edgecolor='black')
        plt.scatter(X[noise_mask, 0], X[noise_mask, 1],
                    c='gray', s=30, alpha=0.5, label='Noise')
        plt.title("Clustering Result")
        plt.legend()
        
        plt.tight_layout()
        plt.show()

# 示例用法
if __name__ == "__main__":
    from sklearn.datasets import make_blobs
    # 设置随机种子保证可重复性
    np.random.seed(42)

    # 生成三个主聚类
    centers = [[1, 1], [5, 5], [9, 1]]  # 等边三角形布局
    cluster_std = [0.5, 0.6, 0.4]  # 各簇标准差

    # 生成聚类数据 (900个点)
    X_clusters, y_clusters = make_blobs(
        n_samples=900,
        centers=centers,
        cluster_std=cluster_std,
        n_features=2,
        random_state=42
    )

    # 生成噪声数据 (100个点)
    noise_min, noise_max = (-2, -2), (12, 8)  # 噪声覆盖更大范围
    X_noise = np.random.uniform(
        low=noise_min,
        high=noise_max,
        size=(100, 2)
    )

    # 合并数据集
    X = np.vstack([X_clusters, X_noise])
    # y = np.hstack([y_clusters, np.full(100, -1)])  # 噪声标记为-1

    # # 可视化
    # plt.figure(figsize=(10, 6))
    # plt.scatter(X[:, 0], X[:, 1], c=y, cmap='viridis', s=30, 
    #             edgecolors='k', alpha=0.7)
    # plt.scatter(np.array(centers)[:,0], np.array(centers)[:,1], 
    #         marker='*', s=300, c='red', edgecolors='white')
    # plt.title("Synthetic Dataset with 3 Clusters and Noise")
    # plt.xlabel("X-axis")
    # plt.ylabel("Y-axis")
    # plt.grid(True, alpha=0.3)
    # plt.show()
    model = DensityPeakCluster(n_clusters=3)
    model.fit(X)
    model.plot_results(X)  # 观察噪声识别和簇分离效果